///
/// Copyright (C) 2016, Dependable Systems Laboratory, EPFL
/// Copyright (C) 2015-2016, Cyberhaven
/// All rights reserved.
///
/// Licensed under the Cyberhaven Research License Agreement.
///

#ifndef S2E_PLUGINS_RECORDSEARCHER_H
#define S2E_PLUGINS_RECORDSEARCHER_H

#include <s2e/record_searcher/commands.h>

#include <s2e/CorePlugin.h>
#include <s2e/Plugin.h>
#include <s2e/Plugins/Core/BaseInstructions.h>
#include <s2e/Plugins/OSMonitors/Support/ModuleExecutionDetector.h>
#include <s2e/Plugins/Searchers/MultiSearcher.h>
#include <s2e/Plugins/StaticAnalysis/ControlFlowGraph.h>
#include <s2e/Plugins/Support/KeyValueStore.h>
#include <s2e/S2EExecutionState.h>
#include <s2e/Synchronization.h>
#include <s2e/Utils.h>

#include <klee/Searcher.h>

#include <llvm/ADT/DenseSet.h>

#include <memory>
#include <queue>
#include <random>
#include <string>
#include <s2e/Plugins/Support/Libwrapper.h>
#include "CUPASearcher.h"

namespace s2e {
namespace plugins {
namespace records {

class RecordSearcher : public Plugin, public klee::Searcher, public IPluginInvoker {
    S2E_PLUGIN

public:

    RecordSearcher(S2E *s2e) : Plugin(s2e) {
    }
    ~RecordSearcher();
    void initialize();

    virtual klee::ExecutionState &selectState();
    virtual void update(klee::ExecutionState *current, const klee::StateSet &addedStates,
                        const klee::StateSet &removedStates);

    virtual bool empty();

private:
    int m_nodeID;

    typedef std::set<S2EExecutionState *> States;
    MultiSearcher *m_multiSearcher;
    CUPASearcher *m_cupa;

    States m_states;

    States m_recordStates;

    States m_masterStates;

    S2EExecutionState *m_cachedState;
    /// Stores the pointer to state 0
    S2EExecutionState *m_initialState;

    std::string m_statefilename;

    /// Location of the record files
    std::string m_recordDirectory;

    /// Enables or disables record scheduling
    bool m_enableRecords;

    /// Stores a backup of the fetched records in the s2e-last/records folder
    bool m_backupRecords;

    /// Location where records will be backed up
    std::string m_recordBackupDirectory;

    /// How many record states are allowed concurrently
    unsigned m_maxRecordStates;

    bool m_parallelRecords;

    //only allow set to 0
    int m_poolSize;

    bool m_hasRecord;

    bool m_isFirstForkDone;

    Libwrapper* m_libwrapper;

    void switchToCUPA();
    void switchToRecordSearcher();

    void onStateSplit(klee::StateSet &parent, klee::StateSet &child);

    void onProcessForkComplete(bool isChild);

    void onStateFork(S2EExecutionState *oldState, const std::vector<S2EExecutionState *> &newStates,
                     const std::vector<klee::ref<klee::Expr>> &);

    void onStateKill(S2EExecutionState *state);

    void backupRecord(const std::string &recordFilePath);
    void fetchNewRecords();
    bool scheduleNextRecord();
    void handleGetRecordFile(S2EExecutionState *state, S2E_RECORDSEARCHER_COMMAND &cmd);

    void handleOpcodeInvocation(S2EExecutionState *state, uint64_t guestDataPtr, uint64_t guestDataSize);

public:
    Libwrapper *getLibwrapper() const {
        return m_libwrapper;
    }
    ///
    /// \brief enableRecords controls whether the searcher schedules records
    ///
    /// When records are disabled, the searcher will still keep polling
    /// for them and notifying clients. Clients can choose to enable records
    /// based on these notifications.
    ///
    /// \param enable is true to use records, false otherwise
    ///
    void enableRecords(bool enable);

    ///
    /// \brief isAvailable returns availability of the searcher on the current S2E instance
    ///
    /// Only the S2E instance that has the state 0 has the record searcher available. Other
    /// instances cannot fetch records because they don't have state 0.
    ///
    /// \return the availability status
    ///
    bool isAvailable() const {
        return m_initialState != nullptr;
    }

};

} // namespace records
} // namespace plugins
} // namespace s2e

#endif // S2E_PLUGINS_RECORDSEARCHER_H
